#pragma once
#include "yolact.h"
#include <fstream>

Yolact::Yolact() {
   this->mTargetSize = 550;
   this->mConfidenceThreshold = 0.5;
   this->mNMSThreshold = 0.5;
   this->mKeepTopK = 200;
   this->mMaskHeight = 138;
   this->mMaskWidth = 138;
   float MEANS[3] = { 123.68, 116.78, 103.94 };
   this->mMEANS = new float[3];
   for (unsigned short i = 0; i < 3; i++) {
      this->mMEANS[i] = MEANS[i];
   }
   float STD[3] = { 58.40, 57.12, 57.38 };
   this->mSTD = new float[3];
   for (unsigned short i = 0; i < 3; i++) {
      this->mSTD[i] = STD[i];
   }
   int convWs[5] = { 69, 35, 18, 9, 5 };
   this->mConvWs = new int[5];
   for (unsigned short i = 0; i < 5; i++) {
      this->mConvWs[i] = convWs[i];
   }
   int convHs[5] = { 69, 35, 18, 9, 5 };
   this->mConvHs = new int[5];
   for (unsigned short i = 0; i < 5; i++) {
      this->mConvHs[i] = convHs[i];
   }
   float aspectRatio[3] = { 1.f, 0.5f, 2.f };
   this->mAspectRatio = new float[3];
   for (unsigned short i = 0; i < 3; i++) {
      this->mAspectRatio[i] = aspectRatio[i];
   }
   float scales[5] = { 24.f, 48.f, 96.f, 192.f, 384.f };
   this->mScales = new float[5];
   for (unsigned short i = 0; i < 5; i++) {
      this->mScales[i] = scales[i];
   }
   float var[4] = { 0.1f, 0.1f, 0.2f, 0.2f };
   this->mVar = new float[4];
   for (unsigned short i = 0; i < 4; i++) {
      this->mVar[i] = var[i];
   }
   this->mNumPriors = 0;
   int p = 0;
   for (p = 0; p < 5; p++) {
      this->mNumPriors += this->mConvWs[p] * this->mConvHs[p] * 3;
   }
   this->mPriorbox = new float[4 * this->mNumPriors];
   const char* classNames[81] = {"background",
      "person", "bicycle", "car", "motorcycle", "airplane", "bus",
      "train", "truck", "boat", "traffic light", "fire hydrant",
      "stop sign", "parking meter", "bench", "bird", "cat", "dog",
      "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe",
      "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
      "skis", "snowboard", "sports ball", "kite", "baseball bat",
      "baseball glove", "skateboard", "surfboard", "tennis racket",
      "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl",
      "banana", "apple", "sandwich", "orange", "broccoli", "carrot",
      "hot dog", "pizza", "donut", "cake", "chair", "couch",
      "potted plant", "bed", "dining table", "toilet", "tv", "laptop",
      "mouse", "remote", "keyboard", "cell phone", "microwave", "oven",
      "toaster", "sink", "refrigerator", "book", "clock", "vase",
      "scissors", "teddy bear", "hair drier", "toothbrush"};
   for (unsigned short i = 0; i < 81; i++) {
      this->mClassNames.push_back(classNames[i]);
   }
}

Yolact::~Yolact() {
   if (this->mMEANS) {
      delete[] this->mMEANS;
   }
   if (this->mSTD) {
      delete[] this->mSTD;
   }
   if (this->mConvWs) {
      delete[] this->mConvWs;
   }
   if (this->mConvHs) {
      delete[] this->mConvHs;
   }
   if (this->mAspectRatio) {
      delete[] this->mAspectRatio;
   }
   if (this->mScales) {
      delete[] this->mScales;
   }
   if (this->mVar) {
      delete[] this->mVar;
   }
   if (this->mPriorbox) {
      delete[] this->mPriorbox;
   }
}


bool Yolact::init(float confidenceThreshold, float NMSThreshold, const char* modelPath, PlatForm platForm, const int keepTopK) {
   confidenceThreshold = (confidenceThreshold < 0.1) ? 0.1 : confidenceThreshold;
   confidenceThreshold = (confidenceThreshold > 0.9) ? 0.9 : confidenceThreshold;
   NMSThreshold = (NMSThreshold < 0.1) ? 0.1 : NMSThreshold;
   NMSThreshold = (NMSThreshold > 0.9) ? 0.9 : NMSThreshold;
   std::ifstream modelFile(modelPath);
   if (!modelFile.good()) {
      return false;
   }
   this->mConfidenceThreshold = confidenceThreshold;
   this->mNMSThreshold = NMSThreshold;
   this->mKeepTopK = keepTopK;
   this->mNet = cv::dnn::readNet(modelPath);
   this->mNet.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
   switch (platForm){
   case PlatForm_CPU:
      this->mNet.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
      break;
   case PlatForm_GPU_INTEL_Integrated:
      this->mNet.setPreferableTarget(cv::dnn::DNN_TARGET_OPENCL_FP16);
      break;
   case PlatForm_GPU_NVDIA_Integrated:
      this->mNet.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
      break;
   case PlatForm_GPU_AMD_Integrated:
      this->mNet.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
      break;
   case PlatForm_GPU_INTEL_Discrate:
      this->mNet.setPreferableTarget(cv::dnn::DNN_TARGET_OPENCL_FP16);
      break;
   case PlatForm_GPU_NVDIA_Discrate:
      this->mNet.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
      break;
   case PlatForm_GPU_AMD_Discrate:
      this->mNet.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
      break;
   case PlatForm_GPU_CUDA:
      this->mNet.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA_FP16);
      break;
   default:
      this->mNet.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
      break;
   }
   //generate priorbox
   float* pb = this->mPriorbox;
   for (int p = 0; p < 5; p++) {
      int convW = this->mConvWs[p];
      int convH = this->mConvHs[p];
      float scale = this->mScales[p];
      for (int i = 0; i < convH; i++) {
         for (int j = 0; j < convW; j++) {
            //+0.5 because priors are in center-size notation
            float cx = (j + 0.5f) / convW;
            float cy = (i + 0.5f) / convH;
            for (int k = 0; k < 3; k++) {
               float ar = this->mAspectRatio[k];
               ar = sqrt(ar);
               float w = scale * ar / this->mTargetSize;
               float h = scale / ar / this->mTargetSize;
               //This is for backward compatability with a bug where I made everything square by accident
               //cfg.backbone.use_square_anchors:
               h = w;
               pb[0] = cx;
               pb[1] = cy;
               pb[2] = w;
               pb[3] = h;
               pb += 4;
            }
         }
      }
   }
}


void Yolact::normalize(cv::Mat& img) {
   img.convertTo(img, CV_32F);
   int i = 0, j = 0;
   for (i = 0; i < img.rows; i++) {
      float* pdata = (float*)(img.data + i * img.step);
      for (j = 0; j < img.cols; j++) {
         pdata[0] = (pdata[0] - this->mMEANS[0]) / this->mSTD[0];
         pdata[1] = (pdata[1] - this->mMEANS[1]) / this->mSTD[1];
         pdata[2] = (pdata[2] - this->mMEANS[2]) / this->mSTD[2];
         pdata += 3;
      }
   }
}



void Yolact::sigmoid(cv::Mat& out, int length) {
   float* pdata = (float*)(out.data);
   int i = 0;
   for (i = 0; i < length; i++) {
      pdata[i] = 1.0 / (1 + expf(-pdata[i]));
   }
}


void Yolact::detect(cv::Mat& srcimg) {
   int imgW = srcimg.cols;
   int imgH = srcimg.rows;
   cv::Mat img;
   cv::resize(srcimg, img, cv::Size(this->mTargetSize, this->mTargetSize), cv::INTER_CUBIC);
   cv::cvtColor(img, img, cv::COLOR_BGR2RGB);
   this->normalize(img);
   cv::Mat blob = cv::dnn::blobFromImage(img);
   this->mNet.setInput(blob);
   std::vector<cv::Mat> outs;
   this->mNet.forward(outs, this->mNet.getUnconnectedOutLayersNames());

   //generate proposals
   std::vector<int> classIds;
   std::vector<float> confidences;
   std::vector<cv::Rect> boxes;
   std::vector<int> maskIds;
   const int numClass = outs[1].cols;
   for (int i = 0; i < this->mNumPriors; i++) {
      cv::Mat scores = outs[1].row(i).colRange(1, numClass);
      cv::Point classIdPoint;
      double score;
      cv::minMaxLoc(scores, 0, &score, 0, &classIdPoint);
      if (score > this->mConfidenceThreshold) {
         const float* loc = (float*)outs[0].data + i * 4;
         const float* pb = this->mPriorbox + i * 4;
         float pb_cx = pb[0];
         float pb_cy = pb[1];
         float pb_w = pb[2];
         float pb_h = pb[3];

         float bbox_cx = this->mVar[0] * loc[0] * pb_w + pb_cx;
         float bbox_cy = this->mVar[1] * loc[1] * pb_h + pb_cy;
         float bbox_w = (float)(exp(this->mVar[2] * loc[2]) * pb_w);
         float bbox_h = (float)(exp(this->mVar[3] * loc[3]) * pb_h);
         float obj_x1 = bbox_cx - bbox_w * 0.5f;
         float obj_y1 = bbox_cy - bbox_h * 0.5f;
         float obj_x2 = bbox_cx + bbox_w * 0.5f;
         float obj_y2 = bbox_cy + bbox_h * 0.5f;

         //clip
         obj_x1 = std::max(std::min(obj_x1 * imgW, (float)(imgW - 1)), 0.f);
         obj_y1 = std::max(std::min(obj_y1 * imgH, (float)(imgH - 1)), 0.f);
         obj_x2 = std::max(std::min(obj_x2 * imgW, (float)(imgW - 1)), 0.f);
         obj_y2 = std::max(std::min(obj_y2 * imgH, (float)(imgH - 1)), 0.f);

         classIds.push_back(classIdPoint.x);
         confidences.push_back(score);
         boxes.push_back(cv::Rect((int)obj_x1, (int)obj_y1, (int)(obj_x2 - obj_x1 + 1), (int)(obj_y2 - obj_y1 + 1)));
         maskIds.push_back(i);
      }
   }
}